import torch.nn as nn


class NET(nn.Module):

    def __init__(self, STATES_NUM, ACTION_NUM):
        super(NET, self).__init__()
        self.l1 = nn.Linear(STATES_NUM, 50)
        self.l1.weight.data.normal_(0, 0.1)
        self.l2 = nn.Linear(50, 50)
        self.l2.weight.data.normal_(0, 0.1)
        self.l3 = nn.Linear(50, 50)
        self.l3.weight.data.normal_(0, 0.1)
        # self.l4 = nn.Linear(20, 20)
        # self.l4.weight.data.normal_(0, 0.1)
        self.out = nn.Linear(50, ACTION_NUM)
        self.out.weight.data.normal_(0, 0.1)

    def forward(self, input_x):
        x = self.l1(input_x)
        x = nn.functional.relu(x)
        x = self.l2(x)
        x = nn.functional.relu(x)
        x = self.l3(x)
        x = nn.functional.relu(x)
        # x = self.l4(x)
        # x = nn.functional.relu(x)
        return self.out(x)
